classdef IrregularityDetector
    properties
        image;
        imageGrad;
        imageDescs;

        samplingStep = 10;
        
        minLocSimilarity = 3e-1;        
        mLocSimSigma = 15;
        locEvalRad = 10;
        locEvalMask;
        locEvalValues;
        
        regionRad = 20;
        patchRad = 3;
        
        progEliminate = 1;
        progEliminateBin = 50;
        
        minDescSimilarity = 3e-1;
        mPatchSimSigma = 1e-2;        
        
        tileSize = 100;
    end
    
    methods (Access = 'public')
        function self = IrregularityDetector(image)
            self.image = image;            
            [ dx dy ] = gradient(self.image);
            self.imageGrad = sqrt(dx.^2 + dy.^2);
            self.imageDescs = self.MakeImageDescs;
            self.locEvalMask = self.MakeLocEvalMask;
            
            % The location of our irregularity.
            ensembleLoc = [ 51 63 ];
            ensemble = self.ObserveEnsemble(ensembleLoc(1), ensembleLoc(2));
            
            self.PlotEnsemble(ensemble);
            
            likelihood = self.EnsembleLikelihood(ensemble);            
            
        end
    end
    
    methods (Access = 'private', Static = true)
        
    end
    
    methods (Access = 'private', Static = false)
        function locEvalMask = MakeLocEvalMask(self)
            locEvalMask = zeros(2 * self.locEvalRad + 1, 2 * self.locEvalRad + 1);
            gridCenter = [ self.locEvalRad + 1, self.locEvalRad + 1 ];
            for row = 1 : 2 * self.locEvalRad + 1
                for col = 1 : 2 * self.locEvalRad + 1
                    diff = [ row col ] - gridCenter;
                    dist = sqrt(diff * diff');
                    if (dist <= self.locEvalRad)
                        locEvalMask(row, col) = 1;
                    end
                end
            end
        end
        
        function ensemble = ObserveEnsemble(self, centerRow, centerCol)
            ensemble.origin = [ centerRow ; centerCol ];
            
            % Set the bounds of the region to be described by the ensemble.
            minRow = max(1, centerRow - self.regionRad) + self.patchRad;
            minCol = max(1, centerCol - self.regionRad) + self.patchRad;
            maxRow = min(size(self.imageGrad, 1), centerRow + self.regionRad) - self.patchRad;
            maxCol = min(size(self.imageGrad, 1), centerCol + self.regionRad) - self.patchRad;
            
            ensemble.bounds.topLeft = [ minRow minCol ];
            ensemble.bounds.bottomRight = [ maxRow maxCol ];

            % Fill the ensemble with descriptors.
            descIndex = 1;
            for iRow = minRow : self.samplingStep : maxRow
                for iCol = minCol : self.samplingStep :maxCol                 
                    desc = self.imageDescs(iRow, iCol, :);
                    ensemble.descs(descIndex).data = desc;
                    ensemble.descs(descIndex).loc = [ iRow - centerRow ; iCol - centerCol ];
                    descIndex = descIndex + 1;
                end
            end
        end
        
        function imageDescs = MakeImageDescs(self)
            patchSize = 2 * self.patchRad + 1;
            imageDescs = zeros(size(self.image, 1), size(self.image, 2), patchSize, patchSize);
            for iRow = self.patchRad + 1 : size(self.image, 1) - self.patchRad;
                for iCol = self.patchRad + 1 : size(self.image, 2) - self.patchRad;
                    imageDescs(iRow, iCol, :) = self.MakeDesc(iRow, iCol);
                end
            end
        end        
        
        function desc = MakeDesc(self, row, col)
            desc = self.imageGrad(row - self.patchRad : row + self.patchRad, col - self.patchRad : col + self.patchRad);
            desc = desc ./ sum(desc(:));
            desc = desc(:);
        end
        
        function likelihood = EnsembleLikelihood(self, queryEnsemble)
             similarity = self.EnsembleDescriptorSimilarity(queryEnsemble);
             [ candidateOrigins descLocs ] = self.EnsembleCandidateLocs(queryEnsemble, similarity);
             likelihood = prod(candidateOrigins, 3);

             if (0)
                 for i = 1 : size(candidateOrigins, 3)
                     imshow(candidateOrigins(:, :, i));
                     pause;
                 end
             end
             
             [ row col ] = (find(likelihood == max(likelihood(:))));

             refEnsemble.origin =  [ row(1) col(1) ];
             for iDesc = 1 : numel(queryEnsemble.descs)
                 [ refEnsemble.descs(iDesc).loc(1) refEnsemble.descs(iDesc).loc(2) ] = ind2sub(size(self.image), descLocs(refEnsemble.origin(1), refEnsemble.origin(2), iDesc));
                 refEnsemble.descs(iDesc).data = self.imageDescs(refEnsemble.descs(iDesc).loc(1), refEnsemble.descs(iDesc).loc(2), :);
                 refEnsemble.descs(iDesc).loc = refEnsemble.descs(iDesc).loc - refEnsemble.origin;
             end
             
             figure;
             self.PlotEnsemble(refEnsemble);
             
        end
        
        function [ origins descLocs ] = EnsembleCandidateLocs(self, queryEnsemble, similarity)
            [tiles] = CreateImageTiles(size(self.image, 2), size(self.image, 1), self.tileSize, self.tileSize, self.regionRad + 1);
            origins = zeros(size(self.image, 1), size(self.image, 2), numel(queryEnsemble.descs));
            descLocs = zeros(size(self.image, 1), size(self.image, 2), numel(queryEnsemble.descs));            
            for iTile = 1 : size(tiles, 1)
                for jTile = 1 : size(tiles, 2)
                    tileRowInd = tiles(iTile, jTile, 2) : tiles(iTile, jTile, 4);
                    tileColInd = tiles(iTile, jTile, 1) : tiles(iTile, jTile, 3);
                    tileSimilarity = similarity(tileRowInd, tileColInd, :);
                    [ curOriginsTile curDescLocsTile ] = EnsembleCandidateLocsTileProcess(self, queryEnsemble, tileSimilarity);
                    originsTile = origins(tileRowInd, tileColInd, :);
                    descLocsTile = descLocs(tileRowInd, tileColInd, :);
                    [ maxInd ] = find(curOriginsTile(:) - originsTile(:) > 0);
                    originsTile(maxInd) = curOriginsTile(maxInd);                    

                    [ rowInd colInd descInd ] = ind2sub(size(curDescLocsTile), curDescLocsTile(maxInd));
                    rowInd = rowInd + tiles(iTile, jTile, 2) - 1;
                    colInd = colInd + tiles(iTile, jTile, 1) - 1;
                    descLocsTile(maxInd) = sub2ind(size(descLocs), rowInd, colInd, descInd);

                    origins(tileRowInd, tileColInd, :) = originsTile;
                    descLocs(tileRowInd, tileColInd, :) = descLocsTile;                    
                end
            end
        end
        
        % Make a heat map for each descriptor depicting the most likely
        % origin locations. Also, output the location of the descriptor
        % pointing to each location at the heat map.        
        function [ origins descLocs ] = EnsembleCandidateLocsTileProcess(self, queryEnsemble, similarity)            
            imageSize = [ size(similarity, 1) size(similarity, 2) ];
            origins = zeros(imageSize(1), imageSize(2), numel(queryEnsemble.descs));
            descLocs = zeros(imageSize(1), imageSize(2), numel(queryEnsemble.descs));            
            
            % Get the nonzero indices of the mask.
            [ nzMaskRow nzMaskCol ] = find(self.locEvalMask);
            % Center the mask at (0,0).
            nzMaskRow = nzMaskRow - self.locEvalRad - 1;
            nzMaskCol = nzMaskCol - self.locEvalRad - 1;              
            
            validOrigins = ones(imageSize);
           
            for iQueryDesc = 1 : numel(queryEnsemble.descs)
                % Find the nonzero image entries for the current descriptor.
                [ row col ] = find(similarity(:, :, iQueryDesc));                
                % Create the ensemble origin shift from the current
                % descriptor's location.
                ensembleShift = -repmat([ queryEnsemble.descs(iQueryDesc).loc ]', numel(row) * numel(nzMaskRow), 1);
                % Create the masks for each image location.
                masks = repmat([ nzMaskRow nzMaskCol ], numel(row), 1);
                % Calculate the mask locations.
                rowColMatInd = repmat(1 : numel(row) ,numel(nzMaskRow), 1);
                rowColMatInd = rowColMatInd(:);
                rowColMat = [ row(rowColMatInd) col(rowColMatInd) ];
                maskLocs = masks + rowColMat + ensembleShift;
                clear masks rowColMatInd ensembleShift 
                % Crop the masks against image bounds.
                maskLocs(:, 1) = min(max(maskLocs(:, 1), 1), imageSize(1));
                maskLocs(:, 2) = min(max(maskLocs(:, 2), 1), imageSize(2));
                maskLocMat = sub2ind(imageSize, maskLocs(:, 1), maskLocs(:, 2));
                % Remove invalid origins.
                validOriginInd = find(validOrigins);
                validInd = ismember(maskLocMat, validOriginInd);
                maskLocMat(~validInd) = [];
                rowColMat(~validInd, :) = []; % take care to remove all bad references to row/col
                maskLocs(~validInd, :) = [];
                % Calculate the posterior origin location probabilities for
                % each descriptor location.
                queryDescLoc = rowColMat - maskLocs;
                clear maskRowInd maskColInd;
                locPosterior = self.GaussianKernelDist(queryDescLoc, queryEnsemble.descs(iQueryDesc).loc', self.mLocSimSigma, 0);
                clear queryDescLoc;
                % Multiply the loc posteriors by their priors.
                curSimilarity = similarity(:, :, iQueryDesc);
                locProb = locPosterior .* curSimilarity(sub2ind(imageSize, rowColMat(:, 1), rowColMat(:, 2))); 
                % Find the max value at each location. To do this
                % effeciently, we sort the mask locations, find the first
                % and last occurence of each location, and thus we easily
                % access the indices in the mask locations matrix that
                % correspond to each image location.
                [ sortedMaskLocs sortedMaskLocsInd ] = sort(maskLocMat);
                [ uniqueMaskLocs uniqueMaskLocsInd ] = unique(sortedMaskLocs);
                uniqueMaskLocsInd = [ 0 ; uniqueMaskLocsInd ]; %#ok
                curDescOrigins = zeros(prod(imageSize), 1);
                curDescLocs = zeros(prod(imageSize), 1);
                maxLocInd = zeros(numel(uniqueMaskLocs), 1);
                % Iterating over the set of possible ensemble origins in
                % the image.
                for originInd = 1 : numel(uniqueMaskLocs)
                    locMaskInd = sortedMaskLocsInd(uniqueMaskLocsInd(originInd) + 1 : uniqueMaskLocsInd(originInd + 1));
                    [ curDescOrigins(uniqueMaskLocs(originInd)) curMaxInd ] = max(locProb(locMaskInd));                    
                    maxLocInd(originInd) = locMaskInd(curMaxInd);
                end                
                curDescLocs(uniqueMaskLocs) = sub2ind(imageSize, rowColMat( maxLocInd, 1), rowColMat( maxLocInd, 2));
                origins(:, :, iQueryDesc) = reshape(curDescOrigins, imageSize(1), imageSize(2));
                descLocs(:, :, iQueryDesc) = reshape(curDescLocs, imageSize(1), imageSize(2));                                
                
                if (self.progEliminate)
                    flatOrigins = origins(:, :, iQueryDesc);
                    flatOrigins = flatOrigins(:);
                    [ dummy sortInd ] = sort(flatOrigins, 'descend');
                    validOrigins = zeros(imageSize);
                    [ toLeaveRow toLeaveCol ] = ind2sub(imageSize, sortInd(1 : self.progEliminateBin));
                    validOrigins(toLeaveRow, toLeaveCol) = 1;
                end
            end
        end      
        
        function similarity = EnsembleDescriptorSimilarity(self, queryEnsemble)
            similarity = zeros(size(self.image, 1), size(self.image, 2), numel(queryEnsemble.descs));
            for iRow = self.patchRad + 1 : size(self.image, 1) - self.patchRad
                for iCol = self.patchRad + 1 : size(self.image, 2) - self.patchRad
                       refDesc = squeeze(self.imageDescs(iRow, iCol, :))';
                       similarity(iRow, iCol, :) = self.GaussianKernelDist(refDesc, squeeze([queryEnsemble.descs.data ]), self.mPatchSimSigma, self.minDescSimilarity);
                end
            end
        end
        
        function dist = GaussianKernelDist(self, x, y, sigma, cutoff)
            Xt = permute(x, [1 3 2]);
            Yt = permute(y, [3 1 2]);
            euclidDist = sum(abs( Xt( :, ones(1, size(y, 1)), : )- Yt( ones(1, size(x, 1)), :, : ) ).^2, 3);
            dist = exp(-euclidDist / ( 2 * sigma));
            dist(dist < cutoff) = 0;
        end
        
        function PlotEnsemble(self, ensemble)
            imshow(self.image);
            if (isfield(ensemble, 'bounds'))
                PlotRect(ensemble.bounds.topLeft, ensemble.bounds.bottomRight, 'g');
            end

            hold on;
            plot(ensemble.origin(2), ensemble.origin(1), 'ro');
            
            for iDesc = 1 : numel(ensemble.descs);
                PlotRect([ ensemble.origin(1) + ensemble.descs(iDesc).loc(1) - self.patchRad, ensemble.origin(2) + ensemble.descs(iDesc).loc(2) - self.patchRad ], ...
                         [ ensemble.origin(1) + ensemble.descs(iDesc).loc(1) + self.patchRad, ensemble.origin(2) + ensemble.descs(iDesc).loc(2) + self.patchRad ], 'c');
            end
        end
    end
end   
%                    if (DEBUG)
%                        imshow(similarity(:, :, iQueryDesc))
%                         hold on
%                         plot(col(iDescLoc), row(iDescLoc), '+')
%                         hold on
%                         plot(curMaskCol, curMaskRow, 'r+')
%                    end